import os, sys
import numpy as np
import re, math
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from multiprocessing import cpu_count
import threadpool

CAP = 60000


def thread_pool(enter_func,param_list):
  pool = threadpool.ThreadPool(cpu_count()*2)
  requests = threadpool.makeRequests(enter_func, param_list)
  [pool.putRequest(req) for req in requests]
  pool.wait()

class data_set_proveider:
    def __init__(self, top_dir):
        self.top_dir = top_dir
        self.pts_fts = []
        self.data_num = []
        self.label = []
        self.load_label()
        self.load_data()

    @staticmethod
    def readFile(filepath):
        with open(filepath, "r") as f:
            while (True):
                yield f.readline().strip()

    @staticmethod
    def filter_batch(pts_fts_, label_):
        label_ = np.expand_dims(label_, 1)
        pts_fts_label = np.concatenate((pts_fts_, label_), 1)
        rm_index = []
        idx = 0
        for pfl in pts_fts_label:
            if math.sqrt(pfl[0] * pfl[0] + pfl[1] * pfl[1] + pfl[2] * pfl[2]) >= 35:
                rm_index.append(idx)
            idx += 1
        return np.delete(pts_fts_label, rm_index, 0)

    def load_data(self):
        cache_pts = []
        cache_fts = []

        all_inte = os.listdir(os.path.join(self.top_dir, "intensity"))
        i = 0
        for file in all_inte:
            sys.stdout.write('\r>> loading inte %d/%d ' % (
                i + 1, len(all_inte)))
            i += 1
            sys.stdout.flush()
            fts_of_file = []
            idx = 0
            for z in data_set_proveider.readFile(os.path.join(self.top_dir, "intensity", file)):
                if idx > CAP:
                    break
                if z.strip() == "":
                    break
                fts_of_file.append(float(z))
                idx += 1
            while idx < CAP:
                fts_of_file.append(float(0))
                idx += 1
            if cache_fts == []:
                cache_fts = fts_of_file
            else:
                cache_fts = np.vstack((fts_of_file, cache_fts))

        all_pts_file = os.listdir(os.path.join(self.top_dir, "pts"))
        i = 0
        for file in all_pts_file:
            sys.stdout.write('\r>> loading pts %d/%d ' % (
                i + 1, len(all_pts_file)))
            i +=1
            pts_of_file = []
            file_pts_num = 0
            idx = 0
            for z in data_set_proveider.readFile(os.path.join(self.top_dir, "pts", file)):
                if idx > CAP:
                    break
                if z.strip() == "":
                    break
                try:
                    node_pts = []
                    li = [float(j) for j in re.split(',', z)]
                    node_pts.append(li[0])
                    node_pts.append(li[1])
                    node_pts.append(li[2])
                    file_pts_num += 1
                    if pts_of_file == []:
                        pts_of_file = node_pts
                    else:
                        pts_of_file = np.vstack((pts_of_file, node_pts))
                except Exception as ess:
                    break
                idx += 1
            self.data_num.append(file_pts_num)
            while idx < CAP:
                pts_of_file = np.vstack((pts_of_file, [0, 0, 0]))
                idx += 1
            pts_of_file = np.expand_dims(pts_of_file, 0)
            if cache_pts == []:
                cache_pts = pts_of_file
            else:
                cache_pts = np.concatenate((pts_of_file, cache_pts), 0)

        cache_fts = np.expand_dims(cache_fts, 2)
        self.pts_fts = np.concatenate((cache_pts, cache_fts), 2)

    def load_label(self):
        all_file = os.listdir(os.path.join(self.top_dir, "category"))
        i = 0
        for file in all_file:
            sys.stdout.write('\r>> loading label %d/%d ' % (
                i + 1, len(all_file)))
            i += 1
            sys.stdout.flush()

            label_of_file = []
            idx = 0
            for z in data_set_proveider.readFile(os.path.join(self.top_dir, "category", file)):
                if idx > CAP:
                    break
                if z.strip() == "":
                    break
                label_of_file.append(int(z))
                idx += 1
            while idx < CAP:
                label_of_file.append(0)
                idx += 1
            if self.label == []:
                self.label = label_of_file
            else:
                self.label = np.vstack((label_of_file, self.label))

    def get_pts_fts_label(self):
        return (self.pts_fts, np.array(self.data_num), self.label)


def getcolor(var):
    return {
        0: 'c',  # DontCare      青色
        # 0: 'k',
        1: 'b',  # cyclist       蓝色
        2: 'g',  # tricycle      绿色
        3: 'r',  # smallMot      红色
        4: 'k',  # bigMot        黑色
        5: 'm',  # pedestrian    洋红
        6: 'y',  # crowds        黄色
        7: 'w',  # unknown       白色
    }.get(var, 'error')


def case1():
    '''
    data_train:(12137, 2974, 3)  data_num_train :(12137,) label_train :(12137, 2974)
    data_val:(1870, 2974, 3)  data_num_val:(1870,)  label_val:(1870, 2974)
    '''
    provider = data_set_proveider("./data_3d_pts_lit/training")
    data_train, data_num_train, label_train = provider.get_pts_fts_label()
    data_val, data_num_val, label_val = provider.get_pts_fts_label()

    batch_index = 1
    new_color = []

    new_pts = provider.filter_batch(data_train[batch_index], label_train[batch_index])
    for c in new_pts:
        new_color.append(getcolor(c[3]))

    new_pts = new_pts.transpose()
    fig = plt.figure(dpi=120)
    ax = fig.add_subplot(111, projection='3d')
    plt.title('point cloud')
    # ax.scatter(XYZ[0], XYZ[1], XYZ[2], c=color, marker='.', s=2, linewidth=0, alpha=1, cmap='spectral')
    ax.scatter(new_pts[0], new_pts[1], new_pts[2], c=new_color, marker='.', s=2, linewidth=0, alpha=1, cmap='spectral')

    # ax.set_facecolor((0,0,0))
    ax.axis('scaled')
    # ax.xaxis.set_visible(False)
    # ax.yaxis.set_visible(False)
    ax.set_xlabel('X Label')
    ax.set_ylabel('Y Label')
    ax.set_zlabel('Z Label')
    plt.show()


def case2():
    provider = data_set_proveider("/home/leo/Downloads/training_data_set/auto_drive/3dpoints/training")
    data_train, data_num_train, label_train = provider.get_pts_fts_label()
    pass


# if __name__ == '__main__':
#     case1()
